#include<cstdlib>
#include<iostream>

#include<boost/shared_ptr.hpp>

#include<formic/lattice_ed/lattice_ed.h>
#include<formic/lattice_ed/input.h>
#include<formic/lattice_ed/indexing.h>
#include<formic/lattice_ed/ham_actor_heis.h>
#include<formic/lattice_ed/ham_actor_hubb.h>
#include<formic/all_data.h>
#include<formic/command/command.h>
#include<formic/output/output.h>
#include<formic/random/random.h>
#include<formic/numeric/numeric.h>
#include<formic/timing/timing.h>
#include<formic/lanczos.h>
#include<formic/mpi/interface.h>

template<class S> void formic::led::run_lanczos(formic::led::Input & userinp) {

  double lambda;

  int vec_len;

  boost::shared_ptr<formic::MatrixActor<S> > ham;

  if ( userinp.ham_type() == "heisenberg" ) {

    // compute total occupation
    const int t = int( userinp.total_sz() + 0.5 * userinp.nsites() * ( userinp.cps() - 1 ) + 0.5 );

    // compute vector length
    vec_len = formic::n_integer_solutions(userinp("nsites"), t, userinp.cps()-1);

    // create an object used for operating by the Hamiltonian
    ham = boost::shared_ptr<formic::MatrixActor<S> >( new formic::led::HamActorHeis<S>(userinp("nsites"),
                                                                                       t,
                                                                                       userinp("cps"),
                                                                                       userinp("bonds"),
                                                                                       userinp("bond_strengths"),
                                                                                       userinp("triangles"),
                                                                                       userinp("nthread"),
                                                                                       userinp("energy_thresh")) );

  } else if ( userinp.ham_type() == "hubbard" ) {

    if ( userinp.cps() != 2 )
      throw formic::Exception("Hubbard Hamiltonian requires cps = 2");

    vec_len =   formic::n_integer_solutions(userinp("nsites"), userinp("nalpha"), userinp.cps()-1)
              * formic::n_integer_solutions(userinp("nsites"), userinp("nbeta" ), userinp.cps()-1);

    // create an object used for operating by the Hamiltonian
    ham = boost::shared_ptr<formic::MatrixActor<S> >( new formic::led::HamActorHubb<S>(userinp("nsites"),
                                                                                       userinp("nalpha"),
                                                                                       userinp("nbeta"),
                                                                                       userinp("hubbard_u"),
                                                                                       userinp("bonds"),
                                                                                       userinp("bond_strengths"),
                                                                                       userinp("nthread"),
                                                                                       userinp("energy_thresh")) );

  } else {
    throw formic::Exception("unknown Hamiltonian type \"%s\"") % userinp.ham_type();
  }

  const int max_iter = userinp("max_iter");

  // allocate memory for use by the lanczos algorithm
  std::vector<S>            x(vec_len,            formic::zero(S()));
  std::vector<S>           v0(vec_len,            formic::zero(S()));
  std::vector<S>           v1(vec_len,            formic::zero(S()));
  std::vector<S>        t_vec(max_iter*max_iter,  formic::zero(S()));
  std::vector<S>        alpha(max_iter,           formic::zero(S()));
  std::vector<S>         beta(max_iter,           formic::zero(S()));
  std::vector<double>       w(max_iter,           0.0);
  std::vector<S>        xwork(5*max_iter,         formic::zero(S()));
  std::vector<double>   rwork(3*max_iter,         0.0);

  // check whether we should load a previous solution before iterating
  bool restart = userinp("restart");

  // iterate, restarting every so many iterations
  bool converged = false;
  for (int total_iter = 0; total_iter < max_iter && !converged; total_iter += userinp.restart_freq()) {

    // determin how many iterations to perform before restarting or stopping
    const int loop_iter = ( total_iter + userinp.restart_freq() >= max_iter ? max_iter - total_iter : userinp.restart_freq() );

    // run lanczos
    converged = formic::lanczos(userinp.compute_vector(),
                                vec_len,
                                loop_iter,
                                lambda,
                                &x.at(0),
                                &v0.at(0),
                                &v1.at(0),
                                &t_vec.at(0),
                                &alpha.at(0),
                                &beta.at(0),
                                &w.at(0),
                                &xwork.at(0),
                                &rwork.at(0),
                                ham.get(),
                                userinp("energy_thresh"),
                                restart);

    // after a restart, start from the best previous vector
    restart = true;

  }

  // compute triple product for the given triangle
  if ( userinp.ham_type() == "heisenberg" ) {
    formic::led::HamActorHeis<S> * ptr = dynamic_cast<formic::led::HamActorHeis<S> *>(ham.get());
    if (ptr == 0)
      throw formic::Exception("dynamic cast of hamiltonian to heisenberg actor failed");
    ptr->triple_products(vec_len, 1, userinp("tp_tri_num"));
  }

}

int formic::led::main(int argc, char ** argv) {

  formic::mpi::init(argc, argv);

  if (formic::mpi::rank() != 0)
    throw formic::Exception("lattice_ed does not support more than one process.");

  formic::start_timer("lattice_ed total");

  if (argc != 3)
    throw formic::Exception("formic::led::main expects argc = 3  (program_name, input_file, output_file), but you gave argc = %i") % argc;

  // open the output file
  formic::open_output_file(argv[2]);

  // read input options from the input file
  formic::init_input_commands();
  formic::led::Input userinp;
  formic::AllData data(userinp);
  formic::run_text_commands(argv[1], data);

  // print input options
  formic::of << "------------------ Printing Input Options ------------------" << std::endl;
  formic::of << std::endl;
  formic::of << userinp.print_options();
  formic::of << std::endl;
  formic::of << "------------------- End of Input Options -------------------" << std::endl;
  formic::of << std::endl;

  // set random seed
  formic::set_seed(userinp("random_seed"));

  // run lanczos
  if ( userinp.n_tp_vecs() == 1 ) {

    if ( userinp.complex_wfn() )
      formic::led::run_lanczos<std::complex<double> >(userinp);
    else
      formic::led::run_lanczos<double>(userinp);

  // diagonalize triple products
  } else if ( userinp.n_tp_vecs() > 1 ) {

    // compute total occupation
    const int t = int( userinp.total_sz() + 0.5 * userinp.nsites() * ( userinp.cps() - 1 ) + 0.5 );

    const int vec_len = formic::n_integer_solutions(userinp("nsites"), t, userinp.cps()-1);
    formic::led::HamActorHeis<std::complex<double> > hamc(userinp("nsites"),
                                                          t,
                                                          userinp("cps"),
                                                          userinp("bonds"),
                                                          userinp("bond_strengths"),
                                                          userinp("triangles"),
                                                          userinp("nthread"),
                                                          userinp("energy_thresh"));
    hamc.triple_products(vec_len, userinp("n_tp_vecs"), userinp("tp_tri_num"));

  } else {
    throw formic::Exception("user entered %i for n_tp_vecs, which should be >= 1") % userinp.n_tp_vecs();
  }

  formic::stop_timer("lattice_ed total");
  formic::of << formic::print_timers();

  formic::close_output_file();

  formic::mpi::finalize();

  return EXIT_SUCCESS;

}
